{-# LANGUAGE TypeInType #-}
{-# LANGUAGE TypeOperators #-}

-- | Module: Data.Vector.Sized.PartialSort
-- Copyright: (c) 2018 Koz Ross
-- Maintainer: Koz Ross <koz.ross@retro-freedom.nz>
-- Stability: Experimental
-- Portability: Non-portable
--
-- Partial sorts for immutable fixed-length vectors, using heapsort with a
-- quarternary heap. These return partially
-- sorted copies; for versions that work in-place, see
-- 'Data.Vector.Sized.Mutable.PartialSort'.
module Data.Vector.Sized.PartialSort
  ( partialSort
  , partialSortBy
  , partialSortWithin
  , partialSortByWithin
  ) where

import Control.Monad.ST
import Data.Finite
import Data.Functor.Contravariant
import GHC.TypeNats

import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable.Sized as VGMS
import qualified Data.Vector.Generic.Sized as VGS
import qualified Data.Vector.Mutable.Sized.PartialSort as VMSP

-- | Partial sorting. After @partialSort v k@, the resulting vector will have
-- its first @k@ elements in order, with the rest in an unspecified order.
--
-- __Complexity:__ \(O(n + k \log (n))\)
{-# INLINE partialSort #-}
partialSort ::
     (KnownNat n, Ord a, VG.Vector v a)
  => VGS.Vector v n a
  -> Finite (n + 1)
  -> VGS.Vector v n a
partialSort vec i =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSP.partialSort v i
        VGS.unsafeFreeze v)

-- | As 'partialSort, but with a custom 'Comparison'.
{-# INLINE partialSortBy #-}
partialSortBy ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> Finite (n + 1)
  -> VGS.Vector v n a
partialSortBy comp vec i =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSP.partialSortBy comp v i
        VGS.unsafeFreeze v)

-- | Partially sorts within the specified indices. More precisely,
-- @partialSortWithin v i j@ performs a partial sort of a copy of @v@ for all
-- elements of @v@ between @min i j@ (inclusive) and @max i j@ (exclusive). In
-- particular, when @i == j@, no sorting will be performed on the copy. Elements
-- outside of this index range will be in no particular order.
--
-- __Complexity:__ \(O(k \log (k))\), where \(k = |j - i|\)
{-# INLINE partialSortWithin #-}
partialSortWithin ::
     (KnownNat n, VG.Vector v a, Ord a)
  => VGS.Vector v n a
  -> Finite n
  -> Finite n
  -> VGS.Vector v n a
partialSortWithin vec i j =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSP.partialSortWithin v i j
        VGS.unsafeFreeze v)

-- | As 'partialSortWithin', but with a custom 'Comparison'.
{-# INLINE partialSortByWithin #-}
partialSortByWithin ::
     (KnownNat n, VG.Vector v a)
  => Comparison a
  -> VGS.Vector v n a
  -> Finite n
  -> Finite n
  -> VGS.Vector v n a
partialSortByWithin comp vec i j =
  runST
    (do v <- VGMS.unsafeNew
        VGS.copy v vec
        VMSP.partialSortByWithin comp v i j
        VGS.unsafeFreeze v)
